import socket
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from Crypto.Hash import SHA256
import rsa
import os
import time

# 加密文件
def encrypt_file(file_path, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, iv)

    with open(file_path, 'rb') as f:
        file_data = f.read()

    encrypted_data = cipher.encrypt(pad(file_data, AES.block_size))

    encrypted_file_path = file_path + '.enc'
    with open(encrypted_file_path, 'wb') as f:
        f.write(encrypted_data)

    print(f"文件已加密并保存为: {encrypted_file_path}")
    return encrypted_file_path


# 签名文件
def sign_file(file_path, private_key_path, signature_path):
    with open(file_path, 'rb') as f:
        file_data = f.read()

    h = SHA256.new(file_data)

    with open(private_key_path, 'rb') as priv_file:
        priv_key = rsa.PrivateKey.load_pkcs1(priv_file.read())

    signature = rsa.sign(h.digest(), priv_key, 'SHA-256')

    print("文件签名完成。")
    with open(signature_path, 'wb') as sig_file:
        sig_file.write(signature)

    print(f"文件签名已成功生成并写入: {signature_path}")
    return signature


# 客户端设置
def start_client(host, port, file_path, key, iv, private_key_path,signature_path):
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect((host, port))

    encrypted_file_path = encrypt_file(file_path, key, iv)
    signature = sign_file(file_path, private_key_path,signature_path)

    # 发送文件大小
    file_size = os.path.getsize(encrypted_file_path)
    client_socket.send(str(file_size).encode().rjust(16, b'0'))

    # 发送加密的文件内容
    with open(encrypted_file_path, 'rb') as f:
        while (data := f.read(1024)):
            client_socket.send(data)
    time.sleep(0.5)
    # 发送签名大小
    client_socket.send(str(len(signature)).encode().rjust(256, b'0'))

    # 发送签名
    # client_socket.send(signature)
    with open(signature_path, 'rb') as f:
        while (data := f.read(16)):
            client_socket.send(data)

    print("文件和签名传输完成。")
    client_socket.close()


if __name__ == "__main__":
    # 设置 AES 加密的密钥和 IV
    # key = b'1234567890123456'  # 16字节的密钥
    with open('shared_secret_B_path.txt', 'rb') as key_file:
        key = key_file.read()
    key = key + b'\x00' * (16 - len(key))
    iv = b'1234567890123456'  # 16字节的初始化向量
    private_key_path = "private.pem"
    signature_path = 'signature.sig'
    start_client('127.0.0.1', 12345, 'file_to_send.txt', key, iv, private_key_path,signature_path)
